package biz.datalk.industrialland.common.component.filter;

import biz.datalk.commons.utils.IdProvider;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.slf4j.MDC;
import org.springframework.core.task.TaskDecorator;
import org.springframework.web.context.request.RequestAttributes;
import org.springframework.web.context.request.RequestContextHolder;

import java.util.HashMap;
import java.util.Map;

/**
 * @author tarofang@163.com
 * @date 2022-06-17
 */
public class MDCFilterTaskDecorator implements TaskDecorator {
    private static final Logger log = LoggerFactory.getLogger(MDCFilterTaskDecorator.class);

    private boolean transferRequestAttributes = Boolean.TRUE; // 默认为 true

    public MDCFilterTaskDecorator() {
    }

    public MDCFilterTaskDecorator(Boolean transferRequestAttributes) {
        this.transferRequestAttributes = transferRequestAttributes == null ? Boolean.TRUE : transferRequestAttributes;
    }

    @Override
    public Runnable decorate(Runnable runnable) {
        RequestAttributes tmpContext = null;
        if (this.transferRequestAttributes) {
            try {
                tmpContext = RequestContextHolder.currentRequestAttributes();
            } catch (Exception ex) {
                if (log.isTraceEnabled()) {
                    log.trace("{}", ex.getMessage(), ex);
                }
                if (log.isWarnEnabled()) {
                    log.warn("{}", ex.getMessage());
                }
            }
        }

        final RequestAttributes context = tmpContext;

        Map<String, String> tmpPrevious;
        try {
            tmpPrevious = MDC.getCopyOfContextMap();
        } catch (Exception ex) {
            tmpPrevious = null;
            if (log.isTraceEnabled()) {
                log.trace("{}", ex.getMessage(), ex);
            }
            if (log.isWarnEnabled()) {
                log.warn("{}", ex.getMessage());
            }
        }
        final Map<String, String> previous = tmpPrevious;

        return () -> {
            try {
                try {
                    if (context != null) {
                        RequestContextHolder.setRequestAttributes(context);
                    }
                    Map<String, String> map = previous == null ? new HashMap<>() : previous;
                    map.putIfAbsent(MDCFilter.TRACE_ID, IdProvider.createUUIDId());
                    MDC.setContextMap(map);
                } catch (Exception ex) {
                    log.error("ERR: {}", ex.getMessage());
                }
                runnable.run();
            } finally {
                RequestContextHolder.resetRequestAttributes();
                MDC.clear();
            }
        };// end of "return () -> {"
    }
}
